import json
import os

import torch
from tqdm import tqdm

from UnetDatasets import UnetDataset
from UnetModel import UnetModel
from config import train_config as train_config
from config import project_path, dataset_path


# 训练主函数入口
def train():
    print('start')

    # 模型
    model = UnetModel()
    start_epoch = 0
    model.to(train_config['device'])

    start_epoch += 1

    train_dataset = UnetDataset(dataset_path=dataset_path, split="train")

    train_data_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                                    batch_size=train_config['batch_size'],
                                                    shuffle=True)

    test_dataset = UnetDataset(dataset_path=dataset_path, split="test")

    test_data_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                                   batch_size=train_config['test_batch_size'],
                                                   shuffle=False)

    # 优化器
    loss_F = torch.nn.MSELoss()
    loss_F.to(train_config['device'])
    optimizer = torch.optim.Adam(model.parameters(), lr=train_config['lr'])

    min_avg_loss = 1e6
    for epoch in range(start_epoch, train_config['epochs'], 1):
        model.train()
        total_loss = 0.0
        min_loss = 10000
        max_loss = 10
        # 按批次取文件
        for (x, y) in tqdm(train_data_loader):
            img = x.to(train_config['device'])
            label = y.to(train_config['device'])
            pre = model(img)  # 前向传播
            mask = (label != 0)

            # 使用掩码来计算损失，只计算那些 label 不为 0 的位置的损失
            loss = loss_F(pre[mask], label[mask])
            optimizer.zero_grad()  # 因为每次反向传播的时候，变量里面的梯度都要清零
            loss.backward()  # 变量得到了grad
            optimizer.step()  # 更新参数
            total_loss += loss.item()
            if loss < min_loss:
                min_loss = loss
            if loss > max_loss:
                max_loss = loss

        avg_loss = total_loss/len(train_data_loader)

        print(f"epoch:{epoch}, avg loss:{avg_loss:.4f}")
        if epoch % train_config['save_epoch'] == 0:

            avg_loss, avg_pck, avg_precision, avg_recall = test(model=model,
                                                                test_loader=test_data_loader,
                                                                loss_F=loss_F,
                                                                device=train_config['device'],
                                                                msg="test")
            saveModel(model, epoch, avg_loss, avg_pck, avg_precision, avg_recall)
            if avg_loss < min_avg_loss:
                min_avg_loss = avg_loss
                saveModel(model,
                          avg_loss=avg_loss,
                          avg_pck=avg_pck,
                          avg_precision=avg_precision,
                          avg_recall=avg_recall, epoch="min_loss")
                print(f"更新最优模型, 当前epoch:{epoch}")
            # test(model=model, test_loader=train_data_loader,
            # loss_F=loss_F, device=train_config['device'], msg="train")


def saveModel(model, epoch, avg_loss, avg_pck, avg_precision, avg_recall):
    save_dir = os.path.join(project_path, "result", "weights", str(epoch))
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    torch.save(model.state_dict(), os.path.join(save_dir, f"epoch_{epoch}.pkl"))
    with open(os.path.join(save_dir, f"result_{epoch}.json"), 'w', encoding="utf-8") as f:
        json.dump(obj={
            "avg_loss": avg_loss,
            "avg_pck": avg_pck,
            "avg_precision": avg_precision,
            "avg_recall": avg_recall
        }, fp=f)


# 计算欧几里得距离
def euclidean_distance(pred, target):
    return torch.sqrt(torch.sum((pred - target) ** 2, dim=-1))


# 计算 PCK、Precision 和 Recall
def calculate_metrics(predictions, targets, threshold=0.05):
    distances = euclidean_distance(predictions, targets)
    true_positives = (distances < threshold).float()  # 小于阈值的点为TP
    false_negatives = (distances >= threshold).float()  # 大于阈值的点为FN
    total_points = targets.size(0)  # 总的点数

    # 计算 Precision 和 Recall
    TP = true_positives.sum().item()
    FN = false_negatives.sum().item()
    FP = (true_positives == 0).float().sum().item()  # 假设所有预测到的点中，超过阈值的点为FP

    precision = TP / (TP + FP + 1e-10)  # 避免除以0
    recall = TP / (TP + FN + 1e-10)

    return precision, recall


# 计算 PCK（Percentage of Correct Keypoints）
def calculate_pck(predictions, targets, threshold=0.05):
    distances = euclidean_distance(predictions, targets)
    correct_keypoints = distances < threshold  # 小于阈值的点
    pck = correct_keypoints.float().mean().item()  # 正确关键点的平均值
    return pck


# 测试函数
def test(model, test_loader, loss_F, device, threshold=10, msg="train"):
    model.eval()  # 设置模型为评估模式
    total_loss = 0.0
    total_pck = 0.0
    total_precision = 0.0
    total_recall = 0.0
    total_samples = 0

    with torch.no_grad():  # 禁用梯度计算
        for (x, y) in tqdm(test_loader):
            img = x.to(device)
            label = y.to(device)
            pre = model(img)  # 前向传播

            # 使用掩码来计算损失，只计算那些 label 不为 0 的位置的损失
            mask = (label != 0)
            loss = loss_F(pre[mask], label[mask])
            total_loss += loss.item()

            # 计算 PCK、Precision 和 Recall
            pre_reshaped = pre.view(-1, 2)
            label_reshaped = label.view(-1, 2)

            batch_pck = calculate_pck(pre_reshaped, label_reshaped, threshold=threshold)  # PCK
            precision, recall = calculate_metrics(pre_reshaped, label_reshaped,
                                                  threshold=threshold)  # Precision & Recall

            total_pck += batch_pck * x.size(0)
            total_precision += precision * x.size(0)
            total_recall += recall * x.size(0)
            total_samples += x.size(0)

    avg_loss = total_loss / len(test_loader)
    avg_pck = total_pck / total_samples
    avg_precision = total_precision / total_samples
    avg_recall = total_recall / total_samples

    print(f"当前测试：{msg}")
    print(f"平均损失: {avg_loss:.4f}")
    print(f"PCK: {avg_pck:.4f}")
    print(f"Precision: {avg_precision:.4f}")
    print(f"Recall: {avg_recall:.4f}")

    return avg_loss, avg_pck, avg_precision, avg_recall


if __name__ == "__main__":
    # 训练
    train()
